
import numpy as np
import datetime
import pysolid
import pandas as pd
from scipy import interpolate
def get_solid(date_string, ndays=180):
    
    lon, lat = 115.814586,39.117921                 # point of interest in degree, Los Angles, CA
    step_sec = 60 * 60                       # sample spacing in time domain in seconds
    ts = datetime.datetime.strptime(date_string, '%Y-%m-%dT%H:%M:%SZ')  # start date and time
    te = ts+datetime.timedelta(days=ndays)  # end   date and time

    dt_out, tide_e, tide_n, tide_u = pysolid.calc_solid_earth_tides_point(
        lat, lon, ts, te,
        step_sec=step_sec,
        display=False,
        verbose=False,
    )
    dt_out = dt_out-ts
    dt_out = np.array([i.days+i.seconds/24/3600+8/24 for i in dt_out])
    
    return dt_out,tide_u

def get_era5_T(date_string, file = 'loc/era5.T.2325.grib',key='t2m'):
    import xarray as xr
    import datetime
    import pandas as pd

    # 使用cfgrib引擎读取grib文件
    ds = xr.open_dataset(file, engine='cfgrib')
    
    # 访问温度变量（通常为't'或't2m'）
    if key in ds:
        temperature = ds[key]  # 2米高温度
    elif 't' in ds:
        temperature = ds['t']    # 默认温度变量
    else:
        # 如果不确定变量名，可以打印所有变量名
        print("Available variables:", list(ds.variables.keys()))
        raise ValueError("Could not find temperature variable")
    
    
    # 获取时间维度
    time = ds.time
    lat = ds.latitude
    lon = ds.longitude
    
    print("Time dimension:", time)
    print("Latitude dimension:", lat)
    print("Longitude dimension:", lon)
    
    # 将时间转换为datetime对象
    ts = datetime.datetime.strptime(date_string, '%Y-%m-%dT%H:%M:%SZ')
    
    # 提取时间序列数据（这里假设我们要提取特定经纬度点的数据）
    # 以北京附近为例 (latitude=39.9, longitude=116.4)
    target_lat, target_lon = 38.61462333,115.2452038
    
    # 找到最接近的网格点
    lat_idx = abs(lat - target_lat).argmin()
    lon_idx = abs(lon - target_lon).argmin()
    
    # 提取时间序列
    temp_series = temperature.isel(latitude=lat_idx, longitude=lon_idx)
    
    # 转换为pandas DataFrame便于处理
    df = temp_series.to_dataframe().reset_index()
    
    # 处理时间
    t_era5 = []
    temp_values = []
    
    for i, row in df.iterrows():
        # 计算相对于起始时间的时间差（天）
        dti = (pd.to_datetime(row['time']) - pd.to_datetime(ts)).total_seconds()/3600/24 + 8/24
        t_era5.append(dti)
        # 温度值需要从Kelvin转换为摄氏度（如果是Kelvin）
        temp_values.append(row[temperature.name] - 273.15 if temperature.units == 'K' else row[temperature.name])
    
    t_era5 = np.array(t_era5)
    temp_values = np.array(temp_values)
    

    return t_era5, temp_values
# ... existing code ...

def get_sensor_T(date_string, file = './loc/P320T.csv'):
    
    ts = datetime.datetime.strptime(date_string, '%Y-%m-%dT%H:%M:%SZ')  # start date and time

    t_phy = []
    data = pd.DataFrame(pd.read_csv(file))
    # 2024-10-13 08:00:00
    for i in data['datetime']:
        di = datetime.datetime.strptime(i, '%Y-%m-%d %H:%M:%S')
        dti = (di - ts).total_seconds()/3600/24+8/24
        t_phy.append(dti)
    t_phy = np.array(t_phy)

    phy = []
    for i in data['T']:
        phy.append(float(i))
    phy = np.array(phy)
    
    return t_phy,phy
# ... existing code ...

def merge_data(t, data, t_merged, method='interp'):

    if method=='interp':
        f4inerp = interpolate.interp1d(t, data, kind='linear', bounds_error=False,fill_value=data.mean())
        return f4inerp(t_merged)    
    else:
        deltaT_real = []
        te_real = []
        ne = len(t_merged)
        for j, t_merged_j in enumerate(t_merged):
            
            L = t_merged[max(j-1,0)]
            R = t_merged[min(j+1,ne-1)]
            mask_j = (t<=L) & (t>R)
            if mask_j.sum()>0:
                deltaT_real.append(data[mask_j].mean())
                te_real.append(t[mask_j].mean())
        f4inerp = interpolate.interp1d(te_real, deltaT_real, kind='linear', bounds_error=False,fill_value=0)
        dt4merge = f4inerp(t_merged)
        return dt4merge
    